import torch
import torch.nn as nn
import numpy as np
from gns import graph_network
from torch_geometric.nn import radius_graph

class LearnedSimulator(nn.Module):
  """Learned simulator from https://arxiv.org/pdf/2002.09405.pdf."""

  def __init__(
          self,
          #stl_path,
          particle_dimensions: int,
          nnode_in: int,
          nedge_in: int,
          latent_dim: int,
          nmessage_passing_steps: int,
          nmlp_layers: int,
          mlp_hidden_dim: int,
          connectivity_radius: float,
          boundaries: np.ndarray,
          normalization_stats: dict,
          nparticle_types: int,
          particle_type_embedding_size: int,
          boundary_clamp_limit: float = 1.0,
          device="cpu"
  ):
    """Initializes the model.

    Args:
      particle_dimensions: Dimensionality of the problem.
      nnode_in: Number of node inputs.
      nedge_in: Number of edge inputs.
      latent_dim: Size of latent dimension (128)
      nmessage_passing_steps: Number of message passing steps.
      nmlp_layers: Number of hidden layers in the MLP (typically of size 2).
      connectivity_radius: Scalar with the radius of connectivity.
      boundaries: Array of 2-tuples, containing the lower and upper boundaries
        of the cuboid containing the particles along each dimensions, matching
        the dimensionality of the problem.
      normalization_stats: Dictionary with statistics with keys "acceleration"
        and "velocity", containing a named tuple for each with mean and std
        fields, matching the dimensionality of the problem.
      nparticle_types: Number of different particle types.
      particle_type_embedding_size: Embedding size for the particle type.
      boundary_clamp_limit: a factor to enlarge connectivity radius used for computing
        normalized clipped distance in edge feature.
      device: Runtime device (cuda or cpu).
    """
    # self._mesh = trimesh.load(stl_path)
    # self._mesh.apply_scale(0.001)  # 毫米转米
    #
    # # 创建加速查询结构
    # # self._sdf_kdtree = KDTree(self._mesh.vertices)
    # self._sdf_fn = SDF(self._mesh.vertices, self._mesh.faces)

    super(LearnedSimulator, self).__init__()
    self._boundaries = boundaries
    self._connectivity_radius = connectivity_radius
    self._normalization_stats = normalization_stats
    self._nparticle_types = nparticle_types
    self._boundary_clamp_limit = boundary_clamp_limit

    # Particle type embedding has shape (9, 16)
    self._particle_type_embedding = nn.Embedding(
        nparticle_types, particle_type_embedding_size)

    #print(f"simulator ecode particle_type_embeddings: {self._particle_type_embedding}")

    # Initialize the EncodeProcessDecode
    self._encode_process_decode = graph_network.EncodeProcessDecode(
        nnode_in_features=nnode_in,
        nnode_out_features=particle_dimensions,
        nedge_in_features=nedge_in,
        latent_dim=latent_dim,
        nmessage_passing_steps=nmessage_passing_steps,
        nmlp_layers=nmlp_layers,
        mlp_hidden_dim=mlp_hidden_dim)

    self._device = device

  def forward(self):
    """Forward hook runs on class instantiation"""
    pass

  # def _compute_graph_connectivity(
  #         self,
  #         node_features: torch.tensor,
  #         nparticles_per_example: torch.tensor,
  #         radius: float,
  #         add_self_edges: bool = True):
  #   """Generate graph edges to all particles within a threshold radius
  #
  #   Args:
  #     node_features: Node features with shape (nparticles, dim).
  #     nparticles_per_example: Number of particles per example. Default is 2
  #       examples per batch.
  #     radius: Threshold to construct edges to all particles within the radius.
  #     add_self_edges: Boolean flag to include self edge (default: True)
  #   """
  #   # Specify examples id for particles
  #   batch_ids = torch.cat(
  #       [torch.LongTensor([i for _ in range(n)])
  #        for i, n in enumerate(nparticles_per_example)]).to(self._device)
  #
  #   # radius_graph accepts r < radius not r <= radius
  #   # A torch tensor list of source and target nodes with shape (2, nedges)
  #   edge_index = radius_graph(
  #       node_features, r=radius, batch=batch_ids, loop=add_self_edges, max_num_neighbors=128)
  #
  #   # The flow direction when using in combination with message passing is
  #   # "source_to_target"
  #   receivers = edge_index[0, :]
  #   senders = edge_index[1, :]
  #
  #   return receivers, senders
  def _compute_graph_connectivity(
          self,
          node_features: torch.tensor,
          nparticles_per_example: torch.tensor,
          radius: float,
          add_self_edges: bool = True):
      """Generate graph edges to all particles within a threshold radius

      Args:
        node_features: Node features with shape (nparticles, dim).
        nparticles_per_example: Number of particles per example. Default is 2
          examples per batch.
        radius: Threshold to construct edges to all particles within the radius.
        add_self_edges: Boolean flag to include self edge (default: True)
      """
      # 确保nparticles_per_example是整数列表
      if isinstance(nparticles_per_example, torch.Tensor):
          nparticles_per_example = nparticles_per_example.tolist()

      # 如果nparticles_per_example是标量，转换为列表
      if not isinstance(nparticles_per_example, list):
          nparticles_per_example = [nparticles_per_example]

      # 确保批次索引正确构建
      batch_ids = []
      for i, n in enumerate(nparticles_per_example):
          batch_ids.extend([i] * n)

      batch_ids = torch.tensor(batch_ids, dtype=torch.long, device=self._device)
      # print(f"simulator nparticles_per_example: {nparticles_per_example}")
      # print(f"simulator batch_ids: {batch_ids.shape}")
      # print(f"simulator node_features: {node_features.shape}")

      # 确保批次索引长度与粒子数量匹配
      assert len(batch_ids) == node_features.size(0), (
          f"Batch ids length ({len(batch_ids)}) does not match number of particles ({node_features.size(0)})")

      # radius_graph accepts r < radius not r <= radius
      # A torch tensor list of source and target nodes with shape (2, nedges)
      edge_index = radius_graph(
          node_features, r=radius, batch=batch_ids, loop=add_self_edges, max_num_neighbors=128)

      # The flow direction when using in combination with message passing is
      # "source_to_target"
      receivers = edge_index[0, :]
      senders = edge_index[1, :]

      return receivers, senders

  def _encoder_preprocessor(
          self,
          position_sequence: torch.tensor,
          nparticles_per_example: torch.tensor,
          particle_types: torch.tensor,
          material_property: torch.tensor):
    """Extracts important features from the position sequence. Returns a tuple
    of node_features (nparticles, 30), edge_index (nparticles, nparticles), and
    edge_features (nparticles, 3).

    Args:
      position_sequence: A sequence of particle positions. Shape is
        (nparticles, 6, dim). Includes current + last 5 positions
      nparticles_per_example: Number of particles per example. Default is 2
        examples per batch.
      particle_types: Particle types with shape (nparticles).
      material_property: Shape is (nparticles, 6, 8)
    """
    nparticles = position_sequence.shape[0]
    #print(f"simulator  position_sequence: {position_sequence.shape}")
    most_recent_position = position_sequence[:, -1]  # (n_nodes, 2)
    velocity_sequence = time_diff(position_sequence)
    #print(f"simulator  velocity_sequence: {velocity_sequence.shape}")

    #print(f"_encoder_preprocessor  nparticles_per_example: {nparticles_per_example}")

    senders, receivers = self._compute_graph_connectivity(
        most_recent_position, nparticles_per_example, self._connectivity_radius)

    node_features = []

    # Normalized velocity sequence, merging spatial an time axis.
    velocity_stats = self._normalization_stats["velocity"]
    normalized_velocity_sequence = (
        velocity_sequence - velocity_stats['mean']) / velocity_stats['std']
    #print(f"simulator  normalized_velocity_sequence: {normalized_velocity_sequence.shape}")
    flat_velocity_sequence = normalized_velocity_sequence.reshape(nparticles, -1)
    #print(f"simulator  flat_velocity_sequence: {flat_velocity_sequence.shape}")
    # There are 5 previous steps, with dim 2
    # node_features shape (nparticles, 5 * 2 = 10)
    node_features.append(flat_velocity_sequence)

    #print(f"simulator flat_velocity_sequence: {flat_velocity_sequence.shape}")

    # 检查material_property的维度
    if material_property.dim() == 3:
        # 三维张量 (nparticles, 6, 8) -> 取最后一帧
        last_material = material_property[:, -1, :]  # (nparticles, 8)
    else:
        # 二维张量 (nparticles, 8)
        last_material = material_property
    # 提取距离特征（SDF值）
    distance_features = last_material[..., [0, 4]]  # 索引0和4：滚筒和障碍物的SDF值
    # distance_features = last_material[..., 0]  # 索引0和4：滚筒和障碍物的SDF值

    # 提取法向量特征
    normal_features = last_material[..., [1, 2, 3, 5, 6, 7]]  # 索引1,2,3,5,6,7：滚筒和障碍物的法向量
    # normal_features = last_material[..., [1, 2, 3]]  # 索引1,2,3,5,6,7：滚筒和障碍物的法向量

    # 对距离特征除以连接半径进行归一化
    normalized_distance = distance_features / self._connectivity_radius

    # 组合处理后的特征
    processed_features = torch.cat([normalized_distance, normal_features], dim=-1)

    # 添加处理后的边界特征
    node_features.append(processed_features)


    boundaries = torch.tensor(
        self._boundaries, requires_grad=False).float().to(self._device)
    distance_to_lower_boundary = (
            most_recent_position - boundaries[:, 0][None])
    distance_to_upper_boundary = (
            boundaries[:, 1][None] - most_recent_position)
    distance_to_boundaries = torch.cat(
        [distance_to_lower_boundary, distance_to_upper_boundary], dim=1)
    static_boundary_features = torch.clamp(
        distance_to_boundaries / self._connectivity_radius,
        -self._boundary_clamp_limit, self._boundary_clamp_limit)

    node_features.append(static_boundary_features)

    # Particle type
    if self._nparticle_types > 1:
      particle_type_embeddings = self._particle_type_embedding(particle_types)

      #print(f"simulator ecode particle_types: {particle_types}")
      node_features.append(particle_type_embeddings)

    # Collect edge features.
    edge_features = []

    normalized_relative_displacements = (
        most_recent_position[senders, :] -
        most_recent_position[receivers, :]
    ) / self._connectivity_radius

    # Add relative displacement between two particles as an edge feature
    # with shape (nparticles, ndim)
    edge_features.append(normalized_relative_displacements)

    # Add relative distance between 2 particles with shape (nparticles, 1)
    # Edge features has a final shape of (nparticles, ndim + 1)
    normalized_relative_distances = torch.norm(
        normalized_relative_displacements, dim=-1, keepdim=True)
    edge_features.append(normalized_relative_distances)

    return (torch.cat(node_features, dim=-1),
            torch.stack([senders, receivers]),
            torch.cat(edge_features, dim=-1))

  def _decoder_postprocessor(
          self,
          normalized_acceleration: torch.tensor,
          position_sequence: torch.tensor,
          material_property: torch.tensor) -> torch.tensor:
    """ Compute new position based on acceleration and current position.
    The model produces the output in normalized space so we apply inverse
    normalization.

    Args:
      normalized_acceleration: Normalized acceleration (nparticles, dim).
      position_sequence: Position sequence of shape (nparticles, dim).

    Returns:
      torch.tensor: New position of the particles.

    """
    acceleration_stats = self._normalization_stats["acceleration"]
    acceleration = (normalized_acceleration * acceleration_stats['std']) + acceleration_stats['mean']

    # Use an Euler integrator to go from acceleration to position, assuming
    # a dt=1 corresponding to the size of the finite difference.
    most_recent_position = position_sequence[:, -1]
    most_recent_velocity = most_recent_position - position_sequence[:, -2]

    # print(f"simulator most_recent_velocity: {most_recent_velocity.shape}")
    # print(f"simulator acceleration: {acceleration.shape}")

    # TODO: Fix dt
    new_velocity = most_recent_velocity + acceleration  # * dt = 1
    new_position = most_recent_position + new_velocity  # * dt = 1

    return new_position

  def predict_positions(
          self,
          current_positions: torch.tensor,
          nparticles_per_example: torch.tensor,
          particle_types: torch.tensor,
          material_property: torch.tensor) -> torch.tensor:
    """Predict position based on acceleration.

    Args:
      current_positions: Current particle positions (nparticles, dim).
      nparticles_per_example: Number of particles per example. Default is 2
        examples per batch.
      particle_types: Particle types with shape (nparticles).
      material_property: Shape is (nparticles, 8)

    Returns:
      next_positions (torch.tensor): Next position of particles.
    """
    #print(f"predict_positions nparticles_per_example: {nparticles_per_example}")
    if material_property is not None:
        node_features, edge_index, edge_features = self._encoder_preprocessor(
            current_positions, nparticles_per_example, particle_types, material_property)
    else:
        node_features, edge_index, edge_features = self._encoder_preprocessor(
            current_positions, nparticles_per_example, particle_types)
    predicted_normalized_acceleration = self._encode_process_decode(
        node_features, edge_index, edge_features)
    next_positions = self._decoder_postprocessor(
        predicted_normalized_acceleration, current_positions, material_property)

    return next_positions

  def predict_accelerations(
          self,
          next_positions: torch.tensor,
          position_sequence_noise: torch.tensor,
          position_sequence: torch.tensor,
          nparticles_per_example: torch.tensor,
          particle_types: torch.tensor,
          material_property: torch.tensor):
    """Produces normalized and predicted acceleration targets.

    Args:
      next_positions: Tensor of shape (nparticles_in_batch, dim) with the
        positions the model should output given the inputs.
      position_sequence_noise: Tensor of the same shape as `position_sequence`
        with the noise to apply to each particle.
      position_sequence: A sequence of particle positions. Shape is
        (nparticles, 6, dim). Includes current + last 5 positions.
      nparticles_per_example: Number of particles per example. Default is 2
        examples per batch.
      particle_types: Particle types with shape (nparticles).
      material_property: Shape is (nparticles, 6, 4)

    Returns:
      Tensors of shape (nparticles_in_batch, dim) with the predicted and target
        normalized accelerations.

    """

    # Add noise to the input position sequence.
    #print(f"predict_accelerations nparticles_per_example: {nparticles_per_example}")
    noisy_position_sequence = position_sequence + position_sequence_noise
    #print(f"Simulator  particle_types: {particle_types}")
    #print(f"predict_accelerations  nparticles_per_example: {nparticles_per_example}")

    # Perform the forward pass with the noisy position sequence.
    if material_property is not None:
        node_features, edge_index, edge_features = self._encoder_preprocessor(
            noisy_position_sequence, nparticles_per_example, particle_types, material_property)
    else:
        node_features, edge_index, edge_features = self._encoder_preprocessor(
            noisy_position_sequence, nparticles_per_example, particle_types)
    predicted_normalized_acceleration = self._encode_process_decode(
        node_features, edge_index, edge_features)

    # Calculate the target acceleration, using an `adjusted_next_position `that
    # is shifted by the noise in the last input position.
    next_position_adjusted = next_positions + position_sequence_noise[:, -1]
    target_normalized_acceleration = self._inverse_decoder_postprocessor(
        next_position_adjusted, noisy_position_sequence)
    # As a result the inverted Euler update in the `_inverse_decoder` produces:
    # * A target acceleration that does not explicitly correct for the noise in
    #   the input positions, as the `next_position_adjusted` is different
    #   from the true `next_position`.
    # * A target acceleration that exactly corrects noise in the input velocity
    #   since the target next velocity calculated by the inverse Euler update
    #   as `next_position_adjusted - noisy_position_sequence[:,-1]`
    #   matches the ground truth next velocity (noise cancels out).

    return predicted_normalized_acceleration, target_normalized_acceleration

  def _inverse_decoder_postprocessor(
          self,
          next_position: torch.tensor,
          position_sequence: torch.tensor):
    """Inverse of `_decoder_postprocessor`.

    Args:
      next_position: Tensor of shape (nparticles_in_batch, dim) with the
        positions the model should output given the inputs.
      position_sequence: A sequence of particle positions. Shape is
        (nparticles, 6, dim). Includes current + last 5 positions.

    Returns:
      normalized_acceleration (torch.tensor): Normalized acceleration.

    """
    previous_position = position_sequence[:, -1]
    previous_velocity = previous_position - position_sequence[:, -2]
    next_velocity = next_position - previous_position
    acceleration = next_velocity - previous_velocity

    acceleration_stats = self._normalization_stats["acceleration"]
    normalized_acceleration = (
        acceleration - acceleration_stats['mean']) / acceleration_stats['std']
    return normalized_acceleration

  def save(
          self,
          path: str = 'model.pt'):
    """Save model state

    Args:
      path: Model path
    """
    torch.save(self.state_dict(), path)

  def load(
          self,
          path: str):
    """Load model state from file

    Args:
      path: Model path
    """
    self.load_state_dict(torch.load(path, map_location=torch.device('cpu')))


def time_diff(
        position_sequence: torch.tensor) -> torch.tensor:
  """Finite difference between two input position sequence

  Args:
    position_sequence: Input position sequence & shape(nparticles, 6 steps, dim)

  Returns:
    torch.tensor: Velocity sequence
  """
  return position_sequence[:, 1:] - position_sequence[:, :-1]
